;; We need to be able to define operators, each of which can be
;; commutative and/or associative. FIXME: do we need associativity?

(in-package :azala)

(declaim (optimize (debug 3)))

(defvar *operators* nil
  "List of all operators")

(defun commutative-p (operator)
  (get operator 'commutative))

(defun (setf commutative-p) (commutative-p operator)
  (setf (get operator 'commutative) commutative-p))

(defun associative-p (operator)
  (get operator 'associative))

(defun (setf associative-p) (associative-p operator)
  (setf (get operator 'associative) associative-p))

(defun operator (name &key commutative associative)
  (pushnew name *operators*)
  (setf (commutative-p name) commutative)
  (setf (associative-p name) associative))

(operator '+ :commutative t :associative t)
(operator '- :commutative nil :associative nil)
(operator '* :commutative t :associative t)
(operator '/ :commutative nil :associative nil)
(operator 'expt :commutative nil :associative nil)
(operator 'log :commutative nil :associative nil)

;; Normalization: converting expressions to a canonical form to make
;; them easier to deal with. The drawback to this is that your
;; expressions may be made almost unrecognizable. Deal with it.

(defgeneric normalize (expression)
  (:documentation "Normalize an expression"))

(defmethod normalize ((exp number)) exp)
(defmethod normalize ((exp symbol)) exp)

(defmethod normalize ((exp list))
  (mapcar #'normalize
	  (destructuring-bind (operator x y) exp
	    (apply #'normalize-list
		   (if (commutative-p operator)
		       (normalize-commutative-op operator x y)
		       exp)))))

(defun normalize-commutative-op (operator x y)
  (cond ((numberp x) (list operator x y))
	((numberp y) (list operator y x))
	(t (list operator x y))))

(defgeneric normalize-list (operator x y)
  (:documentation "Expanded internal version of normalize"))

(defmethod normalize-list ((operator (eql '/)) x y)
  (normalize `(* ,x ,(simplify-list 'expt y -1))))

;; Fallback method does nothing
(defmethod normalize-list (operator x y)
  (list operator x y))

;; Simplification
;;;;;;;;;;;;;;;;;

(defgeneric simplify-list (operator x y)
  (:documentation "Expanded version of simplify which works only on
lists"))

(defmacro defsimp (operator arguments &body body)
  "Define a simplification rule"
  (let ((operator-symbol (gensym "operator")))
    `(defmethod simplify-list ((,operator-symbol (eql ',operator))
			       ,@arguments)
      ,@body)))

(defgeneric simplify (expression)
  (:documentation "Simplify an expression"))

(defmethod simplify ((exp number)) exp)
(defmethod simplify ((exp symbol)) exp)

(defmethod simplify ((exp list))
  (destructuring-bind (operator x y) (normalize exp)
    (if (and (eql operator 'expt)
	     (eql y 0))
	1
	(normalize (simplify-list operator (simplify x) (simplify y))))))

(defun sum-of-constant-multiples-p (operator x y)
  "Is an expression the sum of constant multiples? Like ax + by."
  (and (eql operator '+)
       (listp x)
       (listp y)
       (eql (first x) '*)
       (eql (first y) '*)
       (numberp (second x))
       (numberp (second y))))

(defun factor-constant (operator x y)
  "Factor out a constant. If possible, ax + by = d(ax/d + bx/d)"
  (destructuring-bind ((op1 a xx) (op2 b yy)) (list x y)
    (declare (ignore op1 op2))
    (if (and (integerp a)
	     (integerp b))
	(let ((d (gcd a b)))
	  (if (= d 1)
	      (list operator x y)
	      (simplify `(* ,d (+ (* ,(/ a d) ,xx)
				(* ,(/ b d) ,yy))))))
	(list operator x y))))

(defun associatively-linearizable-p (operator x y)
  "Can an expression be associatively linearized?"
  (and (associative-p operator)
       (listp x)
       (listp y)
       (eql (first x) operator)
       (eql (first y) operator)))

(defun associatively-linearize (operator x y)
  (destructuring-bind ((op1 a xx) (op2 b yy)) (list x y)
    (declare (ignore op1 op2))
    (simplify `(,operator ,a (,operator ,xx (,operator ,b ,yy))))))

(defmethod simplify-list :around (operator x y)
  ;; If both are numbers, apply the operator to them both
  (cond ((and (numberp x) (numberp y))
	 (funcall (symbol-function operator) x y))
	;; If possible, factor out a constant. This may not always
	;; make things simpler, but that's just too bad.
	((sum-of-constant-multiples-p operator x y)
	 (factor-constant operator x y))
	;; If something can be associatively linearized, do so.
	((associatively-linearizable-p operator x y)
	 (associatively-linearize operator x y))
	;; Otherwise, do this sort of thing if possible for assoc. ops:
	;; (+ 1 (+ 2 x)) => (+ 3 x)
	((associative-p operator)
	 (if (and (and (numberp x) (listp y))
		  (eql (first y) operator)
		  (numberp (second y)))
	     (simplify-list operator
			    (funcall (symbol-function operator)
				     x (second y))
			    (third y))
	     (call-next-method)))
	;; If all else fails, call the next method and let it handle
	;; the rest. This is usually operation-specific.
	(t (call-next-method))))

;; Fallback does nothing
(defmethod simplify-list (operator x y)
  (list operator x y))

(defsimp + ((x (eql 0)) y)
  (declare (ignore x))
  y)

(defsimp - ((x (eql 0)) y)
  (declare (ignore x))
  (list '* -1 y))

(defsimp - (x (y (eql 0)))
  (declare (ignore y))
  x)

(defsimp * ((x (eql 0)) y)
  (declare (ignore x))
  0)

(defsimp * ((x (eql 1)) y)
  (declare (ignore x))
  y)

;; Unnecessary because of ugly special case to work around bug
(defsimp expt (base (power (eql 0)))
  0)

(defsimp expt (base (power (eql 1))) base)

(defsimp expt (base (power number))
  (if (listp base)
      (destructuring-bind (operator x y) base
	(if (eql operator '*)
	    `(* ,(simplify `(expt ,x ,power))
	      ,(simplify `(expt ,y ,power)))
	    (call-next-method)))
      (list 'expt (simplify base) power)))

;; ln e = 1
(defsimp log ((base (eql '%e)) (power (eql '%e)))
  1)

(defsimp expt ((base (eql '%e)) power)
  (if (and (listp power)
	   (eql (first power) 'log)
	   (eql (third power) '%e))
      (second power)
      (call-next-method)))

(eval-when ()
  (* (+ 3 5) 7)
  ;; Associativity:
  (* (* x y) z) = (* x (* y z))
  (* (* x 5) 4) = (* x (* 5 4))

  (+ 1 (+ 2 x)) => (+ (+ 1 2) x)
  (* 3 (+ (* (/ 1 3) x) y))

  ;; Factoring:
  (+ (* 3 x) (+ (* 6 (expt x 2)) (* 9 y)))
  (+ (* 3 x) (* 3 (+ (* 2 (expt x 2)) (* 3 y))))
  ;; To factor (+ (* a x) (* b y)), find the GCD of a and b (call it d),
  ;; and return this:
  ;; (* d (+ (* a/d x) (* b/d y)))
  (+ (* 6 x) (* 9 y))

  ;; We need this sort of thing for associative operators:
  (* (* a b) (* c d)) => (* a (* b (* c d)))
  (+ (+ a b) (+ c d)) => (+ a (+ b (+ c d))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(defmethod normalize :around ((exp list))
  (if (= (length exp) 2)
      exp
      (call-next-method)))

(defmethod simplify :around ((exp list))
  (if (= (length exp) 2)
      (normalize (simplify-unary (first exp)
				 (simplify (second exp))))
      (call-next-method)))

(defgeneric simplify-unary (operator x)
  (:documentation "Specialized simplify function only for unary operators"))

(defmethod simplify-unary (operator x)
  (list operator x))